#!/usr/bin/env python
# coding: utf-8

# In[2]:


pip install pandas-datareader


# In[4]:


pip install mplfinance


# In[11]:


import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import mplfinance as mpf
from matplotlib.dates import date2num

# 设置 Tushare token
ts.set_token('6f4ed3cc824c2d45087664dec52d197dfc411c36f4ed520ed8c35c50')
pro = ts.pro_api()


def fetch_stock_data(tickers, start_date, end_date, freq='D'):
    """
    从 Tushare 获取股票数据
    :param tickers: 股票代码列表，格式如 ['000001.SZ', '600519.SH']
    :param start_date: 开始日期，格式 'YYYYMMDD'
    :param end_date: 结束日期，格式 'YYYYMMDD'
    :param freq: 数据频率，'D' 表示日线，'60' 表示 60 分钟线等
    :return: 包含各股票数据的字典，键为股票代码，值为对应的 DataFrame
    """
    data = {}
    for ticker in tickers:
        try:
            df = pro.daily(ts_code=ticker, start_date=start_date, end_date=end_date) if freq == 'D'                 else pro.minute(ts_code=ticker, start_date=start_date, end_date=end_date)
            df = df.sort_values(by='trade_date')
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df.set_index('trade_date', inplace=True)
            data[ticker] = df
        except Exception as e:
            print(f"Error fetching data for {ticker}: {e}")
    return data


def generate_trading_signals(df):
    """
    生成交易信号，这里使用简单的移动平均线交叉策略
    :param df: 股票数据 DataFrame
    :return: 包含交易信号的 DataFrame
    """
    short_window = 10
    long_window = 30
    df['SMA_short'] = df['close'].rolling(window=short_window).mean()
    df['SMA_long'] = df['close'].rolling(window=long_window).mean()
    df['Signal'] = 0
    df.loc[df['SMA_short'] > df['SMA_long'], 'Signal'] = 1  # 金叉，产生买入信号
    df.loc[df['SMA_short'] < df['SMA_long'], 'Signal'] = -1  # 死叉，产生卖出信号
    df['Position'] = df['Signal'].diff()
    return df


def select_and_rank_stocks(tickers, stock_data):
    """
    简单的选股与排序示例，这里按平均收盘价降序排列
    :param tickers: 股票代码列表
    :param stock_data: 包含各股票数据的字典
    :return: 排序后的股票代码列表
    """
    avg_close_prices = {}
    for ticker in tickers:
        avg_close = stock_data[ticker]['close'].mean()
        avg_close_prices[ticker] = avg_close
    sorted_tickers = sorted(avg_close_prices, key=avg_close_prices.get, reverse=True)
    return sorted_tickers


def backtest(df, initial_capital=100000, stop_loss=0.1, take_profit=0.2):
    """
    回测交易策略，包含止损和止盈
    :param df: 包含交易信号的股票数据 DataFrame
    :param initial_capital: 初始资金
    :param stop_loss: 止损比例
    :param take_profit: 止盈比例
    :return: 投资组合数据 DataFrame
    """
    positions = pd.DataFrame(index=df.index).fillna(0.0)
    positions['Stock'] = df['Signal']

    portfolio = positions.multiply(df['close'], axis=0)
    pos_diff = positions.diff()

    portfolio['holdings'] = (positions.multiply(df['close'], axis=0)).sum(axis=1)
    portfolio['cash'] = initial_capital - (pos_diff.multiply(df['close'], axis=0)).sum(axis=1).cumsum()
    portfolio['total'] = portfolio['cash'] + portfolio['holdings']
    portfolio['returns'] = portfolio['total'].pct_change()

    holding_price = None
    for i in range(len(df)):
        if df['Position'].iloc[i] == 1:
            holding_price = df['close'].iloc[i]
        elif df['Position'].iloc[i] == -1:
            holding_price = None
        elif holding_price is not None:
            current_price = df['close'].iloc[i]
            if (current_price / holding_price - 1) <= -stop_loss or (current_price / holding_price - 1) >= take_profit:
                portfolio['Stock'].iloc[i] = 0
                portfolio['cash'].iloc[i:] = portfolio['total'].iloc[i]
                portfolio['holdings'].iloc[i:] = 0
                holding_price = None

    return portfolio


def visualize(df, portfolio, ticker):
    """
    可视化 K 线图、交易信号和收益率曲线
    :param df: 股票数据 DataFrame
    :param portfolio: 投资组合数据 DataFrame
    :param ticker: 股票代码
    """
    # 绘制 K 线图和交易信号
    df_ohlc = df[['open', 'high', 'low', 'close']]
    df_ohlc.index = pd.DatetimeIndex(df_ohlc.index)
    apds = [
        mpf.make_addplot(df['SMA_short'], color='b', label=f'SMA_{len(df["SMA_short"].dropna())}'),
        mpf.make_addplot(df['SMA_long'], color='r', label=f'SMA_{len(df["SMA_long"].dropna())}')
    ]
    buy_signals = df[df['Position'] == 1]
    sell_signals = df[df['Position'] == -1]

    # 调试信息，检查数据长度
    print(f"Buy signals length: {len(buy_signals)}")
    print(f"Sell signals length: {len(sell_signals)}")

    apds.extend([
        mpf.make_addplot(buy_signals['close'], scatter=True, markersize=100, marker='^', color='g', label='Buy Signal'),
        mpf.make_addplot(sell_signals['close'], scatter=True, markersize=100, marker='v', color='r', label='Sell Signal')
    ])
    try:
        mpf.plot(df_ohlc, type='candle', style='yahoo', addplot=apds, title=f'{ticker} Stock Price with Trading Signals',
                 ylabel='Price', volume=False)
    except ValueError as e:
        print(f"Error plotting K-line and signals: {e}")

    # 绘制收益率曲线
    plt.figure()
    try:
        plt.plot(portfolio['total'])
        plt.title(f'{ticker} Portfolio Value Over Time')
        plt.xlabel('Date')
        plt.ylabel('Portfolio Value')
        plt.show()
    except ValueError as e:
        print(f"Error plotting return curve: {e}")


def compare_returns(portfolios):
    """
    比较多支股票的收益率曲线
    :param portfolios: 包含各股票投资组合数据的字典
    :return:
    """
    plt.figure(figsize=(12, 6))
    for ticker, portfolio in portfolios.items():
        plt.plot(portfolio['total'], label=ticker)
    plt.title('Portfolio Value Comparison')
    plt.xlabel('Date')
    plt.ylabel('Portfolio Value')
    plt.legend()
    plt.show()


if __name__ == "__main__":
    tickers = ['000001.SZ', '600519.SH']
    start_date = '20200101'
    end_date = '20210101'
    freq = 'D'

    stock_data = fetch_stock_data(tickers, start_date, end_date, freq)
    portfolios = {}

    sorted_tickers = select_and_rank_stocks(tickers, stock_data)
    for ticker in sorted_tickers:
        df = stock_data[ticker]
        df = generate_trading_signals(df)
        portfolio = backtest(df)
        portfolios[ticker] = portfolio
        visualize(df, portfolio, ticker)

    compare_returns(portfolios)


# In[ ]:




